import numpy as np
import pymc3 as pm
import matplotlib.pyplot as plt

# 模拟一些简单的一元线性关系的数据
np.random.seed(1)
true_slope = 2.5   # 斜率的真实值
true_intercept = -0.5  # 截距的真实值
sigma = 1         # 观测噪声的标准偏差
size = 200        # 数据样本大小

x_data = np.linspace(0, 10, size)                    # 构造自变量 X
y_data = true_slope * x_data + true_intercept       # 根据真实斜率构造因变量 Y
noise = np.random.randn(size) * sigma               # 随机添加观测误差 N(0,sigma^2)，得到最终 y 值
plt.scatter(x=x_data,y=y_data+noise,label='Data')    # 绘制散点图显示原始数据及其对应的直线趋势
plt.plot(x_data,true_slope*x_data+true_intercept,'r',label="True Line")
plt.legend()
plt.show()

with pm.Model() as linear_model:
    
    # 定义先验分布
    slope     = pm.Normal('slope', mu=0, sd=10)
    intercept = pm.Normal('intercept', mu=0, sd=10)
    error     = pm.HalfNormal('error', sd=1)

    # 回归方程
    estimation = intercept + slope * x_data
    
    likelihood = pm.Normal('likelihood',mu=estimation ,sd=error, observed=(y_data + noise))

    trace = pm.sample(draws=2000,tune=1000)


pm.traceplot(trace,var_names=['slope','intercept'])
plt.show()
print(pm.summary(trace))